package jwt

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"github.com/MicahParks/keyfunc"
	jwtpb "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/http/jwt_authn/v3"
	"os"
	"strings"
	njt_log "tmlake.com/api-gateway/istio-agent/proxy"
	"tmlake.com/api-gateway/istio-agent/xds/httpfilter"
	httpfilterJWT "tmlake.com/api-gateway/istio-agent/xds/httpfilter/jwtauthentication"
	"tmlake.com/api-gateway/security/pkg/nodeagent/util"
)

const (
	JwtProvider_RemoteJwks = "JwtProvider_RemoteJwks"
	JwtProvider_LocalJwks  = "JwtProvider_LocalJwks"
	JWKWithError           = "Error-IstiodFailedToFetchJwksUri"

	KtyEC  = "EC"
	KtyOKP = "OKP"
	//HMAC
	KtyOct = "oct"
	KtyRSA = "RSA"
)

type PublicKey struct {
	Key            []byte
	AlgorithmsType string
	//RSA needs
	KeyFilePath string
}

type JwtProvider struct {
	Name                 string
	Issuer               string
	Audiences            []string
	JwksSourceSpecifier  string
	FromHeaders          []*jwtpb.JwtHeader
	FromParams           []string
	ForwardPayloadHeader string
	Jwks                 string
	Public               PublicKey
	JwksWithError        bool
}
type JWT struct {
	//map key: JwtProvider name
	JwtProvider map[string]JwtProvider
	Rule        httpfilterJWT.RequirementRule
}

func CreateJWTConfg(fcs map[string]httpfilter.FilterConfig) map[string]JWT {
	jwtOfPort := make(map[string]JWT)
	for port, fc := range fcs {
		config, ok := fc.(httpfilterJWT.Config)
		if !ok {
			njt_log.NjetLogger.Errorf("expected a 'jwtauthentication.Config' type but %T was returned", config)
			continue
		}
		jwtProviders := applyJwtProvider(config.JwtProviders)
		jwt := JWT{
			JwtProvider: jwtProviders,
			Rule:        config.Rules[0],
		}
		if _, ok := jwtOfPort[port]; !ok {
			jwtOfPort[port] = jwt
		}
	}

	return jwtOfPort
}

func applyJwtProvider(jwtProviders map[string]*jwtpb.JwtProvider) map[string]JwtProvider {
	jwtps := make(map[string]JwtProvider)
	for name, jwtProvider := range jwtProviders {
		var jwksSourceSpecifier, jwks, algorithmsType, keyFilePath string
		var jwksWithError bool
		var key []byte
		switch jwtProvider.GetJwksSourceSpecifier().(type) {
		case *jwtpb.JwtProvider_RemoteJwks:
			njt_log.NjetLogger.Errorf("unsuppot RemoteJwks")
		case *jwtpb.JwtProvider_LocalJwks:
			jwksSourceSpecifier = JwtProvider_LocalJwks
			jwks = jwtProvider.GetLocalJwks().GetInlineString()
			err := checkJWK(jwks)
			if err != nil {
				njt_log.NjetLogger.Errorf("JwtProvider %q JWKS with error: %q", name, jwks)
				jwksWithError = true
			}
			key, algorithmsType, keyFilePath = jwks2pem(name, jwks)
		default:
			njt_log.NjetLogger.Errorf("unsuppot JwksSource")
		}

		jwtps[name] = JwtProvider{
			Name:                 name,
			Issuer:               jwtProvider.GetIssuer(),
			Audiences:            jwtProvider.GetAudiences(),
			FromHeaders:          jwtProvider.GetFromHeaders(),
			FromParams:           jwtProvider.GetFromParams(),
			ForwardPayloadHeader: jwtProvider.GetForwardPayloadHeader(),
			JwksSourceSpecifier:  jwksSourceSpecifier,
			Jwks:                 jwks,
			JwksWithError:        jwksWithError,
			Public: PublicKey{
				Key:            key,
				AlgorithmsType: algorithmsType,
				KeyFilePath:    keyFilePath,
			},
		}
	}

	return jwtps
}

func checkJWK(jwk string) error {
	if len(jwk) == 0 {
		return fmt.Errorf("jwks len is 0")
	}

	if strings.Contains(jwk, JWKWithError) {
		return fmt.Errorf("jwks with error")
	}
	return nil
}

func jwks2pem(name, jwk string) ([]byte, string, string) {
	njt_log.NjetLogger.Debugf("Provider Name: %q, jwks: %q", name, jwk)
	var jwksJSON = json.RawMessage(jwk)
	// Create the JWKS from the resource at the given URL.
	jwks, err := keyfunc.NewJSON(jwksJSON)
	if err != nil {
		njt_log.NjetLogger.Errorf("creates a new JWKS fail: %s", err)
	}

	var publicKeyMemory []byte
	var algorithmsType, keyFilePath string
	//jwks contain several types of keys
	//map key: algorithms type
	//map value: public key array
	pubKeySet := make(map[string][]interface{})
	keys := jwks.ReadOnlyKeys()
	for _, key := range keys {
		switch key.(type) {
		case *rsa.PublicKey:
			if _, ok := pubKeySet[KtyRSA]; !ok {
				var ks []interface{}
				ks = append(ks, key)
				pubKeySet[KtyRSA] = ks
			} else {
				pubKeySet[KtyRSA] = append(pubKeySet[KtyRSA], key)
			}
		case *ecdsa.PublicKey:
		case ed25519.PublicKey:
		case []byte:
			if _, ok := pubKeySet[KtyOct]; !ok {
				var ks []interface{}
				ks = append(ks, key)
				pubKeySet[KtyOct] = ks
			} else {
				pubKeySet[KtyOct] = append(pubKeySet[KtyOct], key)
			}
		default:
			njt_log.NjetLogger.Errorf("unsupport Signature Algorithms PublicKey")
		}
	}

	//If jwks has multiple signature algorithms, we use the RSA first
	if len(pubKeySet) > 1 {
		if _, ok := pubKeySet[KtyRSA]; ok {
			if len(pubKeySet[KtyRSA]) != 0 {
				publicKeyMemory, algorithmsType, keyFilePath, _ = generateRSAPublickey(name, pubKeySet[KtyRSA][0])
			}
		}
	} else {
		for keyType, k := range pubKeySet {
			switch keyType {
			case KtyRSA:
				if len(k) != 0 {
					publicKeyMemory, algorithmsType, keyFilePath, _ = generateRSAPublickey(name, k[0])
				}
			case KtyOct:
				if len(k) != 0 {
					v, ok := k[0].([]byte)
					if ok {
						publicKeyMemory = v
						algorithmsType = KtyOct
						keyFilePath = ""
					}
				}
			}
		}
	}

	return publicKeyMemory, algorithmsType, keyFilePath
}

func generateRSAPublickey(name string, pubKey interface{}) ([]byte, string, string, error) {
	var publicKeyMemory []byte
	var algorithmsType, keyFilePath string
	var err error
	rsapubKey, ok := pubKey.(*rsa.PublicKey)
	if ok {
		var pubData []byte
		pubData, err = x509.MarshalPKIXPublicKey(rsapubKey)
		if err != nil {
			njt_log.NjetLogger.Errorf("converts a public key to PKIX fail")
			return publicKeyMemory, algorithmsType, keyFilePath, err
		}

		ret := pem.EncodeToMemory(&pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: pubData,
		})

		njt_log.NjetLogger.Debugf("jwks2pem: \n%s", string(ret))

		keyFilePath, err = util.OutputPublicKeyOfJwk2PEMToDir(name, os.Getenv("OUTPUT_CERTS"), ret)
		if err != nil {
			njt_log.NjetLogger.Errorf("error when output the resource(jwks2pem): %v", err)
			return publicKeyMemory, algorithmsType, keyFilePath, err
		} else {
			publicKeyMemory = ret
			algorithmsType = KtyRSA
			njt_log.NjetLogger.Infof("output the resource to %q", keyFilePath)
		}
	}

	return publicKeyMemory, algorithmsType, keyFilePath, err
}
